package com.junmeng.ayuv.utils

import android.graphics.ImageFormat
import androidx.camera.core.ImageProxy

/**
 * ImageProxy扩展工具类，可以从ImageProxy提取出指定格式的数据
 */
object ImageProxyUtil {

    /**
     * 获得y平面数据
     */
    @Throws(RuntimeException::class)
    fun getYPlane(
        imageProxy: ImageProxy,
        outBytes: ByteArray? = null//为null时，则结果从返回值获取；不为null时，表明用户希望将结果放入此变量，此设计的目的是可以将缓存交给外部管理，可避免频繁创建和释放的问题，因此推荐用户使用此方式
    ): ByteArray {
        checkFormat(imageProxy)
        val width = imageProxy.width
        val height = imageProxy.height
        val ySize = width * height

        var innerBytes: ByteArray? = null
        if (outBytes == null) {
            innerBytes = ByteArray(ySize)
        } else {
            if (outBytes.size != ySize) {
                throw RuntimeException("请确保out的大小等于宽和高的乘积")
            }
        }
        handleYPlane(imageProxy, 0, innerBytes, outBytes)
        return innerBytes ?: outBytes!!
    }

    /**
     * 获得i420格式的数据，如yyyy yyyy uu vv
     */
    @Throws(RuntimeException::class)
    fun getI420Datas(
        imageProxy: ImageProxy,
        outBytes: ByteArray? = null//为null时，则结果从返回值获取；不为null时，表明用户希望将结果放入此变量，此设计的目的是可以将缓存交给外部管理，可避免频繁创建和释放的问题，因此推荐用户使用此方式)
    ): ByteArray {
        checkFormat(imageProxy)
        val width = imageProxy.width
        val height = imageProxy.height
        val size = width * height * 3 / 2

        var innerBytes: ByteArray? = null
        if (outBytes == null) {
            innerBytes = ByteArray(size)
        } else {
            if (outBytes.size != size) {
                throw RuntimeException("请确保out的大小等于宽和高乘积的3/2")
            }
        }
        var position = handleYPlane(imageProxy, 0, innerBytes, outBytes)
        position = handleUPlane(imageProxy, innerBytes, position, outBytes)
        handleVPlane(imageProxy, innerBytes, position, outBytes)
        return innerBytes ?: outBytes!!
    }

    /**
     * 获得nv12格式的数据，如yyyy yyyy uv uv
     */
    @Throws(RuntimeException::class)
    fun getNV12Datas(
        imageProxy: ImageProxy,
        outBytes: ByteArray? = null//为null时，则结果从返回值获取；不为null时，表明用户希望将结果放入此变量，此设计的目的是可以将缓存交给外部管理，可避免频繁创建和释放的问题，因此推荐用户使用此方式)
    ): ByteArray {
        checkFormat(imageProxy)
        val width = imageProxy.width
        val height = imageProxy.height
        val size = width * height * 3 / 2

        var innerBytes: ByteArray? = null
        if (outBytes == null) {
            innerBytes = ByteArray(size)
        } else {
            if (outBytes.size != size) {
                throw RuntimeException("请确保out的大小等于宽和高乘积的3/2")
            }
        }
        var position = handleYPlane(imageProxy, 0, innerBytes, outBytes)
        handleUVPlane(imageProxy.planes[1], width, height, position, innerBytes, outBytes)
        handleUVPlane(imageProxy.planes[2], width, height, width * height + 1, innerBytes, outBytes)
        return innerBytes ?: outBytes!!
    }

    /**
     * 获得nv21格式的数据，如yyyy yyyy vu vu
     */
    @Throws(RuntimeException::class)
    fun getNV21Datas(
        imageProxy: ImageProxy,
        outBytes: ByteArray? = null//为null时，则结果从返回值获取；不为null时，表明用户希望将结果放入此变量，此设计的目的是可以将缓存交给外部管理，可避免频繁创建和释放的问题，因此推荐用户使用此方式)
    ): ByteArray {
        checkFormat(imageProxy)
        val width = imageProxy.width
        val height = imageProxy.height
        val size = width * height * 3 / 2

        var innerBytes: ByteArray? = null
        if (outBytes == null) {
            innerBytes = ByteArray(size)
        } else {
            if (outBytes.size != size) {
                throw RuntimeException("请确保out的大小等于宽和高乘积的3/2")
            }
        }
        var position = handleYPlane(imageProxy, 0, innerBytes, outBytes)
        handleUVPlane(imageProxy.planes[2], width, height, position, innerBytes, outBytes)
        handleUVPlane(imageProxy.planes[1], width, height, width * height + 1, innerBytes, outBytes)
        return innerBytes ?: outBytes!!
    }

    @Throws(RuntimeException::class)
    private fun checkFormat(imageProxy: ImageProxy) {
        if (ImageFormat.YUV_420_888 != imageProxy.format) {
            throw RuntimeException("expect YUV_420_888, now = $imageProxy.format")
        }
    }

    private fun handleYPlane(
        imageProxy: ImageProxy,
        startPosition: Int = 0,//将要写入innerBytes或outBytes的起始位置
        innerBytes: ByteArray?,
        outBytes: ByteArray?
    ): Int {
        val width = imageProxy.width
        val height = imageProxy.height

        val yPlane = imageProxy.planes[0]
        val yBuffer = yPlane.buffer
        yBuffer.rewind()

        //yuv中的y的pixelStride固定为1，因此可以忽略，只关注rowStride即可
        val yRowStride = yPlane.rowStride
        var position = startPosition
        if (yRowStride == width) { //rowStride和宽度一致的情况下可直接拷贝，不用逐行拷贝
            yBuffer.position(position)
            innerBytes?.let {
                yBuffer.get(it, position, width * height)//
            } ?: yBuffer.get(outBytes, position, width * height)
            position += (width * height)
        } else {
            for (row in 0 until height) { //逐行拷贝
                yBuffer.position(row * yRowStride)
                innerBytes?.let {
                    yBuffer.get(
                        it,
                        position,
                        width
                    ) //表示从yBuffer当前位置取width个字节追加到innerBytes的position位置之后
                } ?: yBuffer.get(outBytes, position, width)
                position += width
            }
        }
        return position
    }

    private fun handleUPlane(
        imageProxy: ImageProxy,
        innerBytes: ByteArray?,
        position: Int,//将要写入innerBytes或outBytes的起始位置
        outBytes: ByteArray?
    ): Int {
        var startPos = position
        val width = imageProxy.width
        val height = imageProxy.height

        val uPlane = imageProxy.planes[1]
        val uBuffer = uPlane.buffer
        uBuffer.rewind()

        val uRowStride = uPlane.rowStride
        val uPixelStride = uPlane.pixelStride

        for (row in 0 until height / 2) { //逐行拷贝
            for (i in 0 until width step uPixelStride) {
                uBuffer.position(row * uRowStride + i)
                innerBytes?.let {
                    uBuffer.get(it, startPos, 1)
                } ?: uBuffer.get(outBytes, startPos, 1)
                startPos++
            }
        }
        return startPos
    }

    private fun handleVPlane(
        imageProxy: ImageProxy,
        innerBytes: ByteArray?,
        position: Int,//将要写入innerBytes或outBytes的起始位置
        outBytes: ByteArray?
    ): Int {
        var startPos = position
        val width = imageProxy.width
        val height = imageProxy.height

        val vPlane = imageProxy.planes[2]
        val vBuffer = vPlane.buffer
        vBuffer.rewind()

        val vRowStride = vPlane.rowStride
        val vPixelStride = vPlane.pixelStride

        for (row in 0 until height / 2) { //逐行拷贝
            for (i in 0 until width step vPixelStride) {
                vBuffer.position(row * vRowStride + i)
                innerBytes?.let {
                    vBuffer.get(it, startPos, 1)
                } ?: vBuffer.get(outBytes, startPos, 1)
                startPos++
            }
        }
        return startPos
    }

    private fun handleUVPlane(
        planeProxy: ImageProxy.PlaneProxy,
        width: Int,
        height: Int,
        position: Int,//将要写入innerBytes或outBytes的起始位置
        innerBytes: ByteArray?,
        outBytes: ByteArray?
    ) {
        val buffer = planeProxy.buffer
        var startPos = position
        val uRowStride = planeProxy.rowStride
        val uPixelStride = planeProxy.pixelStride

        for (row in 0 until height / 2) { //逐行拷贝
            for (i in 0 until width step uPixelStride) {
                buffer.position(row * uRowStride + i)
                innerBytes?.let {
                    buffer.get(it, startPos, 1) //表示从buffer当前位置取1个字节追加到innerBytes的startPos位置之后
                } ?: buffer.get(outBytes, startPos, 1)
                startPos += uPixelStride
            }
        }
    }


}